/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Original file: https://github.com/kubernetes/autoscaler/tree/cluster-autoscaler-chart-1.0.2/vertical-pod-autoscaler/pkg/recommender/model/aggregate_container_state_test.go
Copyright (c) 2024 China Unicom Digital Technology Co., Ltd.

*/

package aggregate

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	colocationv1 "openfuyao.com/colocation-management/pkg/apis/v1"
)

var (
	testTimestamp = time.Date(2024, time.January, 1, 12, 0, 0, 0, time.UTC)
	testPodID     = PodID{Namespace: "test-namespace", PodName: "test-pod"}
	testContainer = "test-container"
)

func TestNewAggregateContainerState(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	assert.NotNil(t, state)
	assert.NotNil(t, state.AggregateCPUUsage)
	assert.NotNil(t, state.AggregateMemoryPeaks)
	assert.True(t, state.AggregateCPUUsage.IsEmpty())
	assert.True(t, state.AggregateMemoryPeaks.IsEmpty())
	assert.Zero(t, state.TotalSamplesCount)
	assert.True(t, state.CreationTime.After(testTimestamp.Add(-time.Minute)))
}

func TestMergeContainerState(t *testing.T) {
	aggregationsConfig = nil
	state1 := NewAggregateContainerState()
	state2 := NewAggregateContainerState()

	// Add samples to both states
	sampleTime1 := testTimestamp
	sampleTime2 := testTimestamp.Add(time.Hour)

	state1.AddSample(&ContainerUsageSample{
		MeasureStart: sampleTime1,
		Usage:        CPUAmountFromCores(1.0),
		Request:      CPUAmountFromCores(2.0),
		Resource:     ResourceCPU,
	})

	state2.AddSample(&ContainerUsageSample{
		MeasureStart: sampleTime2,
		Usage:        CPUAmountFromCores(2.0),
		Request:      CPUAmountFromCores(3.0),
		Resource:     ResourceCPU,
	})

	state2.AddSample(&ContainerUsageSample{
		MeasureStart: sampleTime2,
		Usage:        MemoryAmountFromBytes(1e9),
		Request:      MemoryAmountFromBytes(2e9),
		Resource:     ResourceMemory,
	})

	// Merge state2 into state1
	state1.MergeContainerState(state2)

	assert.Equal(t, 2, state1.TotalSamplesCount)
	assert.Equal(t, sampleTime1, state1.FirstSampleStart)
	assert.Equal(t, sampleTime2, state1.LastSampleStart)
	assert.False(t, state1.AggregateCPUUsage.IsEmpty())
	assert.False(t, state1.AggregateMemoryPeaks.IsEmpty())
}

func TestContainerStateAddSampleCPU(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	sampleTime := testTimestamp

	sample := &ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        CPUAmountFromCores(1.5),
		Request:      CPUAmountFromCores(2.0),
		Resource:     ResourceCPU,
	}

	state.AddSample(sample)

	assert.Equal(t, 1, state.TotalSamplesCount)
	assert.Equal(t, sampleTime, state.FirstSampleStart)
	assert.Equal(t, sampleTime, state.LastSampleStart)
	assert.False(t, state.AggregateCPUUsage.IsEmpty())
}

func TestContainerStateAddSampleMemory(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	sampleTime := testTimestamp

	sample := &ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        MemoryAmountFromBytes(512 * 1024 * 1024),      // 512MB
		Request:      MemoryAmountFromBytes(1 * 1024 * 1024 * 1024), // 1GB
		Resource:     ResourceMemory,
	}

	state.AddSample(sample)

	assert.Equal(t, 0, state.TotalSamplesCount) // Memory samples don't count towards total
	assert.False(t, state.AggregateMemoryPeaks.IsEmpty())
}

func TestSubtractSampleMemory(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	sampleTime := testTimestamp

	// First add a sample
	addSample := &ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        MemoryAmountFromBytes(512 * 1024 * 1024),
		Request:      MemoryAmountFromBytes(1 * 1024 * 1024 * 1024),
		Resource:     ResourceMemory,
	}
	state.AddSample(addSample)

	// Then subtract the same sample
	subtractSample := &ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        MemoryAmountFromBytes(512 * 1024 * 1024),
		Request:      MemoryAmountFromBytes(1 * 1024 * 1024 * 1024),
		Resource:     ResourceMemory,
	}
	state.SubtractSample(subtractSample)

	assert.True(t, state.AggregateMemoryPeaks.IsEmpty())
}

func TestSaveToCheckpoint(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	sampleTime := testTimestamp

	// Add some samples
	state.AddSample(&ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        CPUAmountFromCores(1.0),
		Request:      CPUAmountFromCores(2.0),
		Resource:     ResourceCPU,
	})

	state.AddSample(&ContainerUsageSample{
		MeasureStart: sampleTime,
		Usage:        MemoryAmountFromBytes(1e9),
		Request:      MemoryAmountFromBytes(2e9),
		Resource:     ResourceMemory,
	})

	checkpoint, err := state.SaveToCheckpoint()
	assert.NoError(t, err)

	assert.Equal(t, 1, checkpoint.TotalSamplesCount)
	assert.Equal(t, sampleTime, checkpoint.FirstSampleStart.Time)
	assert.Equal(t, sampleTime, checkpoint.LastSampleStart.Time)
	assert.Equal(t, SupportedCheckpointVersion, checkpoint.Version)
	assert.NotNil(t, checkpoint.CPUHistogram)
	assert.NotNil(t, checkpoint.MemoryHistogram)
}

func TestLoadFromCheckpoint(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	checkpointTime := testTimestamp

	checkpoint := &colocationv1.ContainerCheckpointStatus{
		Version:           SupportedCheckpointVersion,
		LastUpdateTime:    metav1.NewTime(time.Now()),
		FirstSampleStart:  metav1.NewTime(checkpointTime),
		LastSampleStart:   metav1.NewTime(checkpointTime.Add(time.Hour)),
		TotalSamplesCount: 5,
	}

	err := state.LoadFromCheckpoint(checkpoint)
	assert.NoError(t, err)
	assert.Equal(t, 5, state.TotalSamplesCount)
	assert.Equal(t, checkpointTime, state.FirstSampleStart)
	assert.Equal(t, checkpointTime.Add(time.Hour), state.LastSampleStart)
}

func TestLoadFromCheckpointVersionMismatch(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	checkpoint := &colocationv1.ContainerCheckpointStatus{
		Version: "unsupported-version",
	}

	err := state.LoadFromCheckpoint(checkpoint)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "unsupported checkpoint version")
}

func TestIsExpired(t *testing.T) {
	aggregationsConfig = nil
	cs := NewAggregateContainerState()
	cs.LastSampleStart = testTimestamp
	cs.TotalSamplesCount = 1
	assert.False(t, cs.isExpired(testTimestamp.Add(7*24*time.Hour)))
	assert.True(t, cs.isExpired(testTimestamp.Add(8*24*time.Hour)))

	csEmpty := NewAggregateContainerState()
	csEmpty.TotalSamplesCount = 0
	csEmpty.CreationTime = testTimestamp
	assert.False(t, csEmpty.isExpired(testTimestamp.Add(7*24*time.Hour)))
	assert.True(t, csEmpty.isExpired(testTimestamp.Add(8*24*time.Hour)))
}

func TestIsEmpty(t *testing.T) {
	aggregationsConfig = nil
	state := NewAggregateContainerState()
	assert.True(t, state.isEmpty())

	state.TotalSamplesCount = 1
	assert.False(t, state.isEmpty())
}
